#include "postprocess.h"
#include "utils/nms.h"

int postprocess(float* output1, float* output2, float* output3, std::vector<Bboxcv>& result_obj_nms, const Parameter& para)
{
    float ratio = std::min(para.input_shape[0] * 1.0f / para.img_shape[0], para.input_shape[1] * 1.0f / para.img_shape[1]);
    int shape_resized[] = {int(para.img_shape[0] * ratio) + 0.5, int(para.img_shape[1] * ratio + 0.5)};

    float offset[] = {(para.input_shape[0] - shape_resized[0]) / 2.0f / para.input_shape[0], (para.input_shape[1] - shape_resized[1]) / 2.0f / para.input_shape[1]};
    float scale[] = {para.input_shape[0] * 1.0f / shape_resized[0], para.input_shape[1] * 1.0f / shape_resized[1]};
    
    float* outputs[] = {output1, output2, output3};
    for (int i = 0; i < 3; i++)
    {
        for (int j = 0; j < para.num_anchor * (para.class_head + 5) * para.head_shape[i] * para.head_shape[i]; j++)
        {
            outputs[i][j] = 1 / (1 + std::exp(-outputs[i][j]));
        }
    }
    std::vector<Bboxcv> result_obj;
    for (int index = 0; index < 3; index++)
    {
        int width_head = para.head_shape[index];
        int height_head = para.head_shape[index];
        int anchor_scaler_x = para.input_shape[0] / width_head;
        int anchor_scaler_y = para.input_shape[1] / width_head;
        float* output = outputs[index];
        for (int w = 0; w < width_head; w++)
        {
            for (int h = 0; h < height_head; h++)
            {
                for (int a = 0; a < para.num_anchor; a++)
                {
                    float anchor_x = para.anchor[index * 6 + 2 * a] * 1.0 / anchor_scaler_x;
                    float anchor_y = para.anchor[index * 6 + 2 * a + 1] * 1.0 / anchor_scaler_y;
                    float exit_p = output[a * (para.class_head + 5) * width_head * height_head + 4 * width_head * height_head + w * height_head + h];
                    float max_value = para.conf_thres;
                    int max_index = -1;
                    for (int k = 5; k < para.class_head + 5; k++)
                    {
                        float score = exit_p * output[a *  (para.class_head + 5) * width_head * height_head + k * width_head * height_head + w * height_head + h];
                        if (score > max_value)
                        {
                            max_value = score;
                            max_index = k;
                        }
                    }
                    if (max_index != -1)
                    {
                        float x_obj = output[a * (para.class_head + 5) * width_head * height_head + 0 * width_head * height_head + w * height_head + h];
                        float y_obj = output[a * (para.class_head + 5) * width_head * height_head + 1 * width_head * height_head + w * height_head + h];
                        float w_obj = output[a * (para.class_head + 5) * width_head * height_head + 2 * width_head * height_head + w * height_head + h];
                        float h_obj = output[a * (para.class_head + 5) * width_head * height_head + 3 * width_head * height_head + w * height_head + h];

                        x_obj = (x_obj * 2.0 - 0.5 + h) / width_head;
                        y_obj = (y_obj * 2.0 - 0.5 + w) / height_head;
                        w_obj = ((w_obj * 2) * (w_obj * 2) * anchor_x) / width_head;
                        h_obj = ((h_obj * 2) * (h_obj * 2) * anchor_y) / height_head;

                        float box_y = (y_obj - offset[1]) * scale[1];
                        float box_x = (x_obj - offset[0]) * scale[0];
                        float box_h = (h_obj) * scale[1];
                        float box_w = (w_obj) * scale[0];

                        float top = (box_y - box_h / 2) * para.img_shape[1];
                        float left = (box_x - box_w / 2) * para.img_shape[0];
                        float bottom = (box_y + box_h / 2) * para.img_shape[1];
                        float right = (box_x + box_w / 2) * para.img_shape[0];

                        int top_i = (int)(top > 0 ? top : 0);
                        int left_i = (int)(left > 0 ? left : 0);
                        int bottom_i = (int)(para.img_shape[1] > bottom ? bottom : para.img_shape[1]);
                        int right_i = (int)(para.img_shape[0] > right ? right : para.img_shape[0]);
                        cv::Rect rect(cv::Point2i(right_i, bottom_i), cv::Point2i(left_i, top_i));
                        result_obj.push_back(Bboxcv(rect, max_value, int(max_index - 5)));
                    }
                }
            }
        }
    }
    result_obj_nms;
    nms(result_obj, result_obj_nms, para.nms_thres);
    return 0;
}